#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Python version: 3.6

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import copy
import numpy as np
from torchvision import datasets, transforms
import torch

from utils.sampling import mnist_iid, mnist_noniid, cifar_iid
from utils.options import args_parser
from models.Update import LocalUpdate
from models.Nets import MLP, CNNMnist, CNNCifar
from models.Fed import FedAvg
from models.test import test_img

import sys


def fed_T(L_T, dataset_name, iid = True):
    sys.argv = ['']
    args = args_parser()
    
    
    args.num_users = len(L_T[0])
    args.epochs = len(L_T)
    args.dataset = dataset_name
    args.iid = iid
    
    
    args.device = torch.device('cuda:{}'.format(args.gpu) if torch.cuda.is_available() and args.gpu != -1 else 'cpu')
    
    # load dataset and split users
    if args.dataset == 'mnist':
        trans_mnist = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
        dataset_train = datasets.MNIST('../data/mnist/', train=True, download=True, transform=trans_mnist)
        dataset_test = datasets.MNIST('../data/mnist/', train=False, download=True, transform=trans_mnist)
        # sample users
        if args.iid:
            dict_users = mnist_iid(dataset_train, args.num_users)
        else:
            dict_users = mnist_noniid(dataset_train, args.num_users, L_T)
    elif args.dataset == 'cifar':
        trans_cifar = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
        dataset_train = datasets.CIFAR10('../data/cifar', train=True, download=True, transform=trans_cifar)
        dataset_test = datasets.CIFAR10('../data/cifar', train=False, download=True, transform=trans_cifar)
        if args.iid:
            dict_users = cifar_iid(dataset_train, args.num_users)
        else:
            exit('Error: only consider IID setting in CIFAR10')
    else:
        exit('Error: unrecognized dataset')
    img_size = dataset_train[0][0].shape
    
    # build model
    if args.model == 'cnn' and args.dataset == 'cifar':
        net_glob = CNNCifar(args=args).to(args.device)
    elif args.model == 'cnn' and args.dataset == 'mnist':
        net_glob = CNNMnist(args=args).to(args.device)
    elif args.model == 'mlp':
        len_in = 1
        for x in img_size:
            len_in *= x
        net_glob = MLP(dim_in=len_in, dim_hidden=200, dim_out=args.num_classes).to(args.device)
    else:
        exit('Error: unrecognized model')
    print(net_glob)
    net_glob.train()

    # copy weights
    w_glob = net_glob.state_dict()
    
    
    

    # training
    loss_train = []
    accuracy_train = []
    cv_loss, cv_acc = [], []
    val_loss_pre, counter = 0, 0
    net_best = None
    best_loss = None
    val_acc_list, net_list = [], []

    w_locals = [w_glob for i in range(args.num_users)]
    
    acc_tests = []
    loss_tests = []
        
    
    for Round in range(args.epochs):
        
            loss_locals = []
            accuracy_locals = []
            
            idxs_users = range(args.num_users)
            for idx in idxs_users:
                local = LocalUpdate(args=args, dataset=dataset_train, idxs=np.random.choice(dict_users[idx], L_T[Round][idx]))
                w, loss, accuracy = local.train(net=copy.deepcopy(net_glob).to(args.device))
                w_locals[idx] = copy.deepcopy(w)
                loss_locals.append(copy.deepcopy(loss))
                accuracy_locals.append(accuracy)
            # update global weights

            w_glob = FedAvg(w_locals, L_T[Round]) #
            
            # copy weight to net_glob
            net_glob.load_state_dict(w_glob)

            # print loss
            loss_avg = sum(loss_locals) / len(loss_locals)
            
            
            accuracy_avg = sum(accuracy_locals) / len(accuracy_locals)
            print('Round {:3d}, Average Train loss {:.3f}, Average Train Accuracy {:.3f}'.format(Round, loss_avg, accuracy_avg))
            
            loss_train.append(loss_avg)
            accuracy_train.append(accuracy_avg)

    

            # testing
            net_glob.eval()
            acc_test, loss_test = test_img(net_glob, dataset_test, args)

            print("Testing accuracy: {:.2f}".format(acc_test))
            print("Testing loss: {:.2f}".format(loss_test))
            acc_tests.append(acc_test)
            loss_tests.append(loss_test)

    return loss_train, accuracy_train, acc_tests, loss_tests, args




